import clang.cindex
import os
import json
import argparse

from utils import is_system_function_declaration, LLVM_LIBRARY_PATH, clang_options, collect_include_paths, system_header_paths

def get_symbol_from_rust(symbol_name, rust_symbol_table):
    info = ""
    if symbol_name in rust_symbol_table:
        return symbol_name, rust_symbol_table[symbol_name]
    elif symbol_name.upper() in rust_symbol_table:
        return symbol_name.upper(), rust_symbol_table[symbol_name.upper()]
    elif symbol_name.upper().replace("_", "") in rust_symbol_table:
        return symbol_name.upper().replace("_", ""), rust_symbol_table[symbol_name.upper().replace("_", "")]
    else:
        lower_symbol_name = symbol_name.lower().replace("_", "")
        for rust_symbol_name in rust_symbol_table.keys():
            lower_rust_symbol_name = rust_symbol_name.lower().replace("_", "")
            if lower_symbol_name == lower_rust_symbol_name:
                return rust_symbol_name, rust_symbol_table[rust_symbol_name]
        return symbol_name, "" 

def visit_node_for_usage(node, c_used_symbols, c_symbol_table, 
                               rust_used_symbols, rust_symbol_table):
    if is_system_function_declaration(node):
        return
    if node.kind == clang.cindex.CursorKind.DECL_REF_EXPR:
        # 找到引用的标识符
        symbol_name = node.spelling
        if symbol_name in c_symbol_table:
            c_used_symbols[symbol_name] = c_symbol_table[symbol_name]
            rust_symbol_name, info = get_symbol_from_rust(symbol_name, rust_symbol_table)
            if info:
                rust_used_symbols[rust_symbol_name] = info
    elif node.kind == clang.cindex.CursorKind.TYPE_REF:
        symbol_name = node.spelling
        if "struct " in symbol_name:
            symbol_name = symbol_name.replace("struct ", "")
        elif "enum " in symbol_name:
            symbol_name = symbol_name.replace("enum ", "")
        elif "union " in symbol_name:
            symbol_name = symbol_name.replace("union ", "")

        if symbol_name in c_symbol_table:
            c_used_symbols[symbol_name] = c_symbol_table[symbol_name]
            rust_symbol_name, info = get_symbol_from_rust(symbol_name, rust_symbol_table)
            if info:
                rust_used_symbols[rust_symbol_name] = info
    elif node.kind == clang.cindex.CursorKind.VAR_DECL or \
            node.kind == clang.cindex.CursorKind.PARM_DECL or \
            node.kind == clang.cindex.CursorKind.FUNCTION_DECL or \
            node.kind == clang.cindex.CursorKind.FIELD_DECL or \
            node.kind == clang.cindex.CursorKind.ENUM_CONSTANT_DECL or \
            node.kind == clang.cindex.CursorKind.CALL_EXPR or \
            node.kind == clang.cindex.CursorKind.MEMBER_REF_EXPR or \
            node.kind == clang.cindex.CursorKind.UNEXPOSED_EXPR or \
            node.kind == clang.cindex.CursorKind.BINARY_OPERATOR or \
            node.kind == clang.cindex.CursorKind.UNARY_OPERATOR:
        symbol_name = node.spelling
        if symbol_name in c_symbol_table:
            c_used_symbols[symbol_name] = c_symbol_table[symbol_name]
            rust_symbol_name, info = get_symbol_from_rust(symbol_name, rust_symbol_table)
            if info:
                rust_used_symbols[rust_symbol_name] = info
    else:
        symbol_name = node.spelling
        if symbol_name and symbol_name in c_symbol_table:
            c_used_symbols[symbol_name] = c_symbol_table[symbol_name]
            rust_symbol_name, info = get_symbol_from_rust(symbol_name, rust_symbol_table)
            if info:
                rust_used_symbols[rust_symbol_name] = info


    # 递归访问子节点
    for child in node.get_children():
        if is_system_function_declaration(child):
            continue
        visit_node_for_usage(child, c_used_symbols, c_symbol_table, rust_used_symbols, rust_symbol_table)

def extract_used_symbols(c_symbols, c_symbol_table, 
                                 rust_symbol_table, include_dirs=[]):
    
    # 初始化clang库
    #if not clang.cindex.Config.loaded:
    #    clang.cindex.Config.set_library_path(LLVM_LIBRARY_PATH)  # 修改为libclang路径
    #
    #index = clang.cindex.Index.create()

    # 创建临时文件以供clang解析
    #with open('temp.c', 'w') as temp_file:
    #    temp_file.write(c_code)

    # 构建编译参数，设置包括目录
    #args = ['-I' + include_dir for include_dir in include_dirs]

    # 解析C代码文件
    #translation_unit = index.parse('temp.c', args=args, options=clang_options)
    
    c_used_symbols = dict()
    rust_used_symbols = dict()
    
    if not c_symbols:
        return c_used_symbols, rust_used_symbols
    # 从根节点开始访问AST
    #visit_node_for_usage(translation_unit.cursor, c_used_symbols, c_symbol_table, rust_used_symbols, rust_symbol_table)

    # 删除临时文件
    #os.remove('temp.c')
    for symbol_name in c_symbols:
        if symbol_name in c_symbol_table:
            c_used_symbols[symbol_name] = c_symbol_table[symbol_name]
            rust_symbol_name, info = get_symbol_from_rust(symbol_name, rust_symbol_table)
            if info:
                rust_used_symbols[rust_symbol_name] = info

    return c_used_symbols, rust_used_symbols

def c_format_symbol(symbol_name, symbol_info):
    category = symbol_info['category']
    if category == 'function':
        # 恢复函数声明
        params = ', '.join(f'{pname} {ptype}' for pname, ptype in symbol_info['parameters'])
        return f"{symbol_info['return_type']} {symbol_name}({params});"
    elif category == 'global_var':
        # 恢复变量定义
        init_value = symbol_info['init_value']
        return f"{symbol_info['type']} {symbol_name}{' = ' + init_value if init_value else ''};"
    elif category == 'struct':
        # 恢复结构体定义
        fields = symbol_info.get('members', {})
        field_defs = ';\n  '.join(f'{ftype} {fname}' for fname, ftype in fields.items())
        return f"struct {symbol_name} {{\n  {field_defs};\n}};"
    elif category == 'enum_member':
        # 恢复枚举定义
        enum_name = symbol_info.get('enum', "")
        return f"enum value:{symbol_name} defined in {enum_name}"
    return ""

def rust_format_symbol(symbol_name, symbol_info):
    #print(symbol_info)
    if isinstance(symbol_info, str):
        return ""
    try:
        category = symbol_info['category']
    except Exception as e:
        category = ""
        print(e, symbol_info)
    if category == 'function':
        # 恢复函数声明
        params = ', '.join(f'{pname} {ptype}' for pname, ptype in symbol_info['parameters'])
        params_list = []
        for pname, ptype in symbol_info['parameters']:
            if ptype:
                params_list.append(f"{pname}: {ptype}")
            else:
                params_list.append(f"{pname}")
        params = ", ".join(params_list)
        return_type = ""
        if symbol_info['return_type']:
            return_type = f" -> {symbol_info['return_type']}"
        return f"fn {symbol_name}({params}){return_type};"
    elif category == 'impl method':
        # 恢复函数声明
        params = ', '.join(f'{pname} {ptype}' for pname, ptype in symbol_info['parameters'])
        params_list = []
        for pname, ptype in symbol_info['parameters']:
            if ptype:
                params_list.append(f"{pname}: {ptype}")
            else:
                params_list.append(f"{pname}")
        params = ", ".join(params_list)
        return_type = ""
        if symbol_info['return_type']:
            return_type = f" -> {symbol_info['return_type']}"
        return f"fn {symbol_name}({params}){return_type};"
    elif category == 'global_var':
        # 恢复变量定义
        kind = symbol_info['attr']
        init_value = symbol_info['init_value']
        return f"{kind} {symbol_name}: {symbol_info['type']}{' = ' + init_value if init_value else ''};"
    elif category == 'typedef':
        # 恢复变量定义
        original_type = symbol_info['original']
        return f"type {symbol_name} = {original_type};"

    elif category == 'struct':
        # 恢复结构体定义
        fields = symbol_info.get('members', {})
        field_defs = ',\n '.join(f'{fname}: {ftype}' for fname, ftype in fields.items())
        if field_defs:
            return f"struct {symbol_name} {{\n  {field_defs};\n}};"
        else:
            return f"struct {symbol_name} {{}};"
    elif category == 'enum_member':
        # 恢复枚举定义
        enum_name = symbol_info.get('enum', "")
        return f"enum value:{symbol_name} defined in {enum_name}"
    return ""

def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="从C代码片段中提取使用的符号")
    #parser.add_argument("src_path", type=str, help="c pro path.")
    parser.add_argument("c_symbols", type=str, help="待解析的C symbols")
    parser.add_argument("c_symbol_table_json", type=str, help="C 符号表的JSON字符串")
    parser.add_argument("rust_symbol_table_json", type=str, help="Rust 符号表的JSON字符串")
    args = parser.parse_args()

    # 解析符号表JSON字符串为Python字典
    try:
        c_symbol_table = json.loads(args.c_symbol_table_json)
    except json.JSONDecodeError as e:
        print(f"Fail to parse C symbol table in JSON due to:{e}")
        return
    
    try:
        rust_symbol_table = json.loads(args.rust_symbol_table_json)
    except json.JSONDecodeError as e:
        print(f"Fail to parse Rust symbol table in JSON due to:{e}")
        return

    # 提取代码片段中使用的符号
    #include_dirs = collect_include_paths(args.src_path)
    #include_dirs.extend(system_header_paths)
    c_symbols = args.c_symbols.split(":")
    c_used_symbols, rust_used_symbols = extract_used_symbols(c_symbols, c_symbol_table, rust_symbol_table)

    # 输出使用的符号信息
    if  len(c_used_symbols.keys()) > 0:
        if len(rust_used_symbols.keys()) > 0:
            print("Here are the symbols used in the C code snippet:\n```c")
            for name, details in c_used_symbols.items():
                c_formatted_symbol = c_format_symbol(name, details)
                print(c_formatted_symbol)
            print("```")
            print("The corresponding symbols in the Rust code snippet are:\n```rust")
            for name, details in rust_used_symbols.items():
                rust_formatted_symbol = rust_format_symbol(name, details)
                print(rust_formatted_symbol)
            print("```")
        else:
            print("Here are the symbols used in the C code snippet:\n```c")
            for name, details in c_used_symbols.items():
                c_formatted_symbol = c_format_symbol(name, details)
                print(c_formatted_symbol)
            print("```")
    
# if __name__ == "__main__":
#     main()


def test_code_snippet_symbol():
    import sys
    import json
    
    # 设置测试文件路径
    test_src_path = "test_files"
    
    # 准备测试用的C代码片段
    test_c_code = """
    Point p = {0, 0};
    move_point(&p, 1, 2);
    """
    
    # 准备测试用的C符号表
    c_symbol_table = {
        "Point": {
            "category": "struct",
            "members": {
                "x": "int",
                "y": "int"
            }
        },
        "move_point": {
            "category": "function",
            "return_type": "void",
            "parameters": [
                ["p", "Point*"],
                ["dx", "int"],
                ["dy", "int"]
            ]
        }
    }
    
    # 准备测试用的Rust符号表
    rust_symbol_table = {
        "Point": {
            "category": "struct",
            "members": {
                "x": "i32",
                "y": "i32"
            }
        },
        "move_point": {
            "category": "function",
            "return_type": "",
            "parameters": [
                ["p", "*mut Point"],
                ["dx", "i32"],
                ["dy", "i32"]
            ]
        }
    }
    
    # 将符号表转换为JSON字符串
    c_symbol_json = json.dumps(c_symbol_table)
    rust_symbol_json = json.dumps(rust_symbol_table)
    
    # 模拟命令行参数
    sys.argv = [
        "extract-code-snippet-symbol.py",
        test_src_path,
        test_c_code,
        c_symbol_json,
        rust_symbol_json
    ]
    
    # 调用main函数
    main()

if __name__ == "__main__":
    # 测试请注释下面的main()调用，改成调用test_code_snippet_symbol()
    #test_code_snippet_symbol()
    main()
